import os
import sys

sys.path.append('/data/zhuoran/code/cognlp')

from cognlp.io.processor.et.ontonotes import OntoNotesEtProcessor
from cognlp.io.loader.et.ontonotes import OntoNotesEtLoader
from cognlp.core.dataset import OntoNotesDataset
import torch
torch.cuda.set_device(5)

import torch.nn as nn
import torch.optim as optim
import torch.optim.lr_scheduler as lr_scheduler
from torch.utils.data import RandomSampler
from cognlp.models.et.bert_et import Bert4Et
from cognlp.core.metrics import MultiLabelStrictAccuracyMetric
from cognlp.core.trainer import Trainer, Tester
from cognlp.utils.util import load_json, save_json

# loader = OntoNotesEtLoader()
# data = loader.load_train("../data/et/OntoNotes/data")
# processor = OntoNotesEtProcessor(label_list=loader.get_labels(), path="../data/et/OntoNotes/data")
# train_data = processor.process(data)
# save_json(train_data, "../data/et/OntoNotes/data/train.json")

processor = OntoNotesEtProcessor(path="../data/et/OntoNotes/data")
train_data = load_json("../data/et/OntoNotes/data/train.json")
device = torch.device('cuda')
train_data = OntoNotesDataset(train_data, device)
train_sampler = RandomSampler(train_data)
model = Bert4Et(len(processor.vocabulary))
loss = nn.BCELoss()
optimizer = optim.Adam(model.parameters(), lr=0.000001)
# scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=[3, 4], gamma=0.9)
metric = MultiLabelStrictAccuracyMetric()

# tester = Tester(model, model_path='../data/et/OntoNotes/model/2020-12-02-13:49:31-model.pkl', batch_size=16, sampler=train_sampler,
#                  drop_last=False, num_workers=0, print_every=1000,
#                  dev_data=train_data, metrics=metric, metric_key=None, use_tqdm=True, device=None,
#                  callbacks=None, check_code_level=0, device_ids=[4, 5, 6])
# tester.test()
trainer = Trainer(train_data=train_data, model=model, optimizer=optimizer, loss=loss,
                  batch_size=20, train_sampler=train_sampler, drop_last=False, gradient_accumulation_steps=1,
                  num_workers=5, n_epochs=10, print_every=1000, scheduler=None, dev_sampler=train_sampler,
                  dev_data=train_data, metrics=metric, metric_key=None, validate_steps=1,
                  save_path="../data/et/OntoNotes/model", save_file=None, save_steps=1, use_tqdm=True, device=device,
                  callbacks=None, check_code_level=0, grad_norm=None, device_ids=[5, 6, 7])
trainer.train()
print(1)
